package avicit.bdp.dds.server.master.consumer;

import avicit.bdp.common.service.dto.DataSourceDTO;
import avicit.bdp.dds.dispatch.enums.ExecutionStatus;
import avicit.bdp.dds.dispatch.enums.ResourceType;
import avicit.bdp.dds.dispatch.enums.TaskType;
import avicit.bdp.dds.dispatch.enums.UdfType;
import avicit.bdp.dds.dispatch.model.TaskNode;
import avicit.bdp.dds.dispatch.process.ResourceInfo;
import avicit.bdp.dds.dispatch.task.AbstractParameters;
import avicit.bdp.dds.dispatch.task.datax.DataxParameters;
import avicit.bdp.dds.dispatch.task.procedure.ProcedureParameters;
import avicit.bdp.dds.dispatch.task.sql.SqlParameters;
import avicit.bdp.dds.dispatch.task.sqoop.SqoopParameters;
import avicit.bdp.dds.dispatch.task.sqoop.sources.SourceMysqlParameter;
import avicit.bdp.dds.dispatch.task.sqoop.targets.TargetMysqlParameter;
import avicit.bdp.core.util.json.JSONUtils;
import avicit.bdp.dds.common.Constants;
import avicit.bdp.dds.common.thread.Stopper;
import avicit.bdp.dds.common.utils.EnumUtils;
import avicit.bdp.dds.common.utils.FileUtils;
import avicit.bdp.dds.common.utils.TaskParametersUtils;
import avicit.bdp.dds.dao.entity.ProcessDefinition;
import avicit.bdp.dds.dao.entity.ProcessInstance;
import avicit.bdp.dds.dao.entity.Resource;
import avicit.bdp.dds.dao.entity.TaskInstance;
import avicit.bdp.dds.dao.entity.UdfFunc;
import avicit.bdp.dds.server.builder.TaskExecutionContextBuilder;
import avicit.bdp.dds.server.entity.DataxTaskExecutionContext;
import avicit.bdp.dds.server.entity.ProcedureTaskExecutionContext;
import avicit.bdp.dds.server.entity.SQLTaskExecutionContext;
import avicit.bdp.dds.server.entity.SqoopTaskExecutionContext;
import avicit.bdp.dds.server.entity.TaskExecutionContext;
import avicit.bdp.dds.server.entity.TaskPriority;
import avicit.bdp.dds.server.master.config.MasterConfig;
import avicit.bdp.dds.server.master.dispatch.ExecutorDispatcher;
import avicit.bdp.dds.server.master.dispatch.context.ExecutionContext;
import avicit.bdp.dds.server.master.dispatch.enums.ExecutorType;
import avicit.bdp.dds.server.master.dispatch.exceptions.ExecuteException;
import avicit.bdp.dds.service.process.ProcessService;
import avicit.bdp.dds.service.queue.TaskPriorityQueue;
import com.alibaba.fastjson2.JSONObject;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @author admin
 * TaskUpdateQueue consumer
 */
@Component
public class TaskPriorityQueueConsumer extends Thread {

    /**
     * logger of TaskUpdateQueueConsumer
     */
    private static final Logger logger = LoggerFactory.getLogger(TaskPriorityQueueConsumer.class);

    /**
     * taskUpdateQueue
     */
    private final TaskPriorityQueue taskPriorityQueue;

    /**
     * processService
     */
    private final ProcessService processService;

    /**
     * executor dispatcher
     */
    private final ExecutorDispatcher dispatcher;


    /**
     * master config
     */
    private final MasterConfig masterConfig;

    @Autowired
    public TaskPriorityQueueConsumer(TaskPriorityQueue taskPriorityQueue,
                                     ProcessService processService,
                                     ExecutorDispatcher dispatcher,
                                     MasterConfig masterConfig) {

        this.taskPriorityQueue = taskPriorityQueue;
        this.processService = processService;
        this.dispatcher = dispatcher;
        this.masterConfig = masterConfig;
    }

    @PostConstruct
    public void init() {
        super.setName("TaskUpdateQueueConsumerThread");
        super.start();
    }

    @Override
    public void run() {
        List<String> failedDispatchTasks = new ArrayList<>();
        while (Stopper.isRunning()) {
            try {
                int fetchTaskNum = masterConfig.getMasterDispatchTaskNumber();
                failedDispatchTasks.clear();
                for (int i = 0; i < fetchTaskNum; i++) {
                    if (taskPriorityQueue.size() <= 0) {
                        Thread.sleep(Constants.SLEEP_TIME_MILLIS);
                        continue;
                    }
                    // if not task , blocking here
                    String taskPriorityInfo = taskPriorityQueue.take();
                    TaskPriority taskPriority = TaskPriority.of(taskPriorityInfo);
                    boolean dispatchResult = dispatch(taskPriority.getTaskId(), taskPriority.getProcessInstanceId());
                    if (!dispatchResult) {
                        failedDispatchTasks.add(taskPriorityInfo);
                    }
                }
                for (String dispatchFailedTask : failedDispatchTasks) {
                    taskPriorityQueue.put(dispatchFailedTask);
                }
            } catch (Exception e) {
                logger.error("dispatcher task error", e);
            }
        }
    }


    /**
     * dispatch task
     *
     * @param taskInstanceId taskInstanceId
     * @return result
     */
    private boolean dispatch(String taskInstanceId, String processInstanceId) {
        boolean result = false;
        try {
            ProcessDefinition definition = processService.getProcessDefinitionByProcessInstanceId(processInstanceId);
            TaskExecutionContext context = null;
            //处理sparkflow节点分发
            String type = definition.getDatasourceId();
            if (type != null && (type.startsWith(ProcessDefinition.DATASOURCE_TYPE_SPARK)
                    || type.startsWith(ProcessDefinition.DATASOURCE_TYPE_PHM))) {
                context = getSparkFlowExecutionContext(taskInstanceId, processInstanceId);
                if (sparkFlowIsFinalState(processInstanceId)) {
                    // when task finish, ignore this task, there is no need to dispatch anymore
                    return true;
                }
            } else {
                context = getTaskExecutionContext(taskInstanceId);
                if (taskInstanceIsFinalState(taskInstanceId)) {
                    // when task finish, ignore this task, there is no need to dispatch anymore
                    return true;
                }
            }
            ExecutionContext executionContext = new ExecutionContext(context.toCommand(), ExecutorType.WORKER, context.getWorkerGroup());
            result = dispatcher.dispatch(executionContext);


        } catch (ExecuteException e) {
            logger.error("dispatch error", e);
        }
        return result;
    }


    /**
     * taskInstance is final state
     * success，failure，kill，stop，pause，threadwaiting is final state
     *
     * @param taskInstanceId taskInstanceId
     * @return taskInstance is final state
     */
    public Boolean taskInstanceIsFinalState(String taskInstanceId) {
        TaskInstance taskInstance = processService.findTaskInstanceById(taskInstanceId);
        return ExecutionStatus.of(taskInstance.getState()).typeIsFinished();
    }

    public Boolean sparkFlowIsFinalState(String processInstanceId) {
        ProcessInstance instance = processService.findProcessInstanceById(processInstanceId);
        return ExecutionStatus.of(instance.getState()).typeIsFinished();
    }

    protected TaskExecutionContext getSparkFlowExecutionContext(String taskInstanceId, String processInstanceId) {
        ProcessInstance instance = processService.findProcessInstanceById(processInstanceId);
        ProcessDefinition definition = processService.getProcessDefinitionByProcessInstanceId(processInstanceId);
        TaskExecutionContext taskExecutionContext = TaskExecutionContextBuilder.get().buildProcessInstanceRelatedInfo(instance).create();
        taskExecutionContext.setProcessDefineId(instance.getProcessDefinitionId());
        taskExecutionContext.setTaskType(TaskType.SPARKFLOW.name());
        taskExecutionContext.setTaskJson(instance.getProcessInstanceJson());
        taskExecutionContext.setTaskName(instance.getName());
        taskExecutionContext.setProjectId(definition.getProjectId());
        taskExecutionContext.setTaskInstanceId(taskInstanceId);

        return taskExecutionContext;
    }

    /**
     * get TaskExecutionContext
     *
     * @param taskInstanceId taskInstanceId
     * @return TaskExecutionContext
     */
    protected TaskExecutionContext getTaskExecutionContext(String taskInstanceId) {
        TaskInstance taskInstance = processService.getTaskInstanceDetailByTaskId(taskInstanceId);
        // task type
        TaskType taskType = TaskType.valueOf(taskInstance.getTaskType());
        // task node
        TaskNode taskNode = JSONObject.parseObject(taskInstance.getTaskJson(), TaskNode.class);

        String tenant = processService.getTenantForProcess();
        String queue = processService.getQueueForProcess();

        // verify tenant is null
        if (verifyTenantIsNull(tenant, taskInstance)) {
            processService.changeTaskState(ExecutionStatus.FAILURE.getCode(), taskInstance.getStartTime(),
                    taskInstance.getHost(), null, null, taskInstance.getId());
            return null;
        }
        // set queue for process instance, user-specified queue takes precedence over tenant queue
        String userQueue = processService.queryUserQueueByProcessInstanceId(taskInstance.getProcessInstanceId());
        taskInstance.getProcessInstance().setQueue(StringUtils.isEmpty(userQueue) ? queue : userQueue);
        taskInstance.getProcessInstance().setTenantCode(tenant);
        taskInstance.setExecutePath(getExecLocalPath(taskInstance));
        taskInstance.setResources(getResourceFullNames(taskNode));

        SQLTaskExecutionContext sqlTaskExecutionContext = null;
        DataxTaskExecutionContext dataxTaskExecutionContext = null;
        ProcedureTaskExecutionContext procedureTaskExecutionContext = null;
        SqoopTaskExecutionContext sqoopTaskExecutionContext = null;

        // SQL task
        if (taskType == TaskType.SQL) {
            sqlTaskExecutionContext = new SQLTaskExecutionContext();
            setSQLTaskRelation(sqlTaskExecutionContext, taskNode);
        } else if (taskType == TaskType.DATAX || taskType == TaskType.DATAXINCREMENT) {
            dataxTaskExecutionContext = new DataxTaskExecutionContext();
            setDataxTaskRelation(dataxTaskExecutionContext, taskNode);
        } else if (taskType == TaskType.PROCEDURE) {
            // procedure task
            procedureTaskExecutionContext = new ProcedureTaskExecutionContext();
            setProcedureTaskRelation(procedureTaskExecutionContext, taskNode);
        } else if (taskType == TaskType.SQOOP) {
            sqoopTaskExecutionContext = new SqoopTaskExecutionContext();
            setSqoopTaskRelation(sqoopTaskExecutionContext, taskNode);
        }

        return TaskExecutionContextBuilder.get()
                .buildTaskInstanceRelatedInfo(taskInstance)
                .buildProcessInstanceRelatedInfo(taskInstance.getProcessInstance())
                .buildProcessDefinitionRelatedInfo(taskInstance.getProcessDefine())
                .buildSQLTaskRelatedInfo(sqlTaskExecutionContext)
                .buildDataxTaskRelatedInfo(dataxTaskExecutionContext)
                .buildProcedureTaskRelatedInfo(procedureTaskExecutionContext)
                .buildSqoopTaskRelatedInfo(sqoopTaskExecutionContext)
                .create();
    }

    /**
     * set procedure task relation
     *
     * @param procedureTaskExecutionContext procedureTaskExecutionContext
     * @param taskNode                      taskNode
     */
    private void setProcedureTaskRelation(ProcedureTaskExecutionContext procedureTaskExecutionContext, TaskNode taskNode) {
        ProcedureParameters procedureParameters = JSONObject.parseObject(taskNode.getParams(), ProcedureParameters.class);
        String datasourceId = procedureParameters.getDatasource();
        DataSourceDTO datasource = processService.findDataSourceById(datasourceId);
        procedureTaskExecutionContext.setConnectionParams(datasource.getConnectionParams());
    }

    /**
     * set datax task relation
     *
     * @param dataxTaskExecutionContext dataxTaskExecutionContext
     * @param taskNode                  taskNode
     */
    private void setDataxTaskRelation(DataxTaskExecutionContext dataxTaskExecutionContext, TaskNode taskNode) {
        DataxParameters dataxParameters = JSONObject.parseObject(taskNode.getParams(), DataxParameters.class);

        DataSourceDTO dataSource = processService.findDataSourceById(dataxParameters.getDataSource());
        DataSourceDTO dataTarget = processService.findDataSourceById(dataxParameters.getDataTarget());

        if (dataSource != null) {
            dataxTaskExecutionContext.setDataSourceId(dataxParameters.getDataSource());
            dataxTaskExecutionContext.setSourcetype(dataSource.getDatasourceType());
            dataxTaskExecutionContext.setSourceConnectionParams(dataSource.getConnectionParams());
        }

        if (dataTarget != null) {
            dataxTaskExecutionContext.setDataTargetId(dataxParameters.getDataTarget());
            dataxTaskExecutionContext.setTargetType(dataTarget.getDatasourceType());
            dataxTaskExecutionContext.setTargetConnectionParams(dataTarget.getConnectionParams());
        }
    }


    /**
     * set datax task relation
     *
     * @param sqoopTaskExecutionContext sqoopTaskExecutionContext
     * @param taskNode                  taskNode
     */
    private void setSqoopTaskRelation(SqoopTaskExecutionContext sqoopTaskExecutionContext, TaskNode taskNode) {
        SqoopParameters sqoopParameters = JSONObject.parseObject(taskNode.getParams(), SqoopParameters.class);

        SourceMysqlParameter sourceMysqlParameter = JSONUtils.parseObject(sqoopParameters.getSourceParams(), SourceMysqlParameter.class);
        TargetMysqlParameter targetMysqlParameter = JSONUtils.parseObject(sqoopParameters.getTargetParams(), TargetMysqlParameter.class);

        DataSourceDTO dataSource = processService.findDataSourceById(sourceMysqlParameter.getSrcDatasource());
        DataSourceDTO dataTarget = processService.findDataSourceById(targetMysqlParameter.getTargetDatasource());

        if (dataSource != null) {
            sqoopTaskExecutionContext.setDataSourceId(dataSource.getId());
            sqoopTaskExecutionContext.setSourcetype(dataSource.getDatasourceType());
            sqoopTaskExecutionContext.setSourceConnectionParams(dataSource.getConnectionParams());
        }

        if (dataTarget != null) {
            sqoopTaskExecutionContext.setDataTargetId(dataTarget.getId());
            sqoopTaskExecutionContext.setTargetType(dataTarget.getDatasourceType());
            sqoopTaskExecutionContext.setTargetConnectionParams(dataTarget.getConnectionParams());
        }
    }

    /**
     * set SQL task relation
     *
     * @param sqlTaskExecutionContext sqlTaskExecutionContext
     * @param taskNode                taskNode
     */
    private void setSQLTaskRelation(SQLTaskExecutionContext sqlTaskExecutionContext, TaskNode taskNode) {
        SqlParameters sqlParameters = JSONObject.parseObject(taskNode.getParams(), SqlParameters.class);
        String datasourceId = sqlParameters.getDatasource();
        DataSourceDTO datasource = processService.findDataSourceById(datasourceId);
        sqlTaskExecutionContext.setConnectionParams(datasource.getConnectionParams());

        // whether udf type
        boolean udfTypeFlag = EnumUtils.isValidEnum(UdfType.class, sqlParameters.getType())
                && StringUtils.isNotEmpty(sqlParameters.getUdfs());

        if (udfTypeFlag) {
            String[] udfFunIds = sqlParameters.getUdfs().split(Constants.COMMA);
            List<UdfFunc> udfFuncList = processService.queryUdfFunListByIds(udfFunIds);
            Map<UdfFunc, String> udfFuncMap = new HashMap<>(16);
            for (UdfFunc udfFunc : udfFuncList) {
                String tenantCode = processService.queryTenantCodeByResName(udfFunc.getFuncName(), ResourceType.UDF);
                udfFuncMap.put(udfFunc, tenantCode);
            }

            sqlTaskExecutionContext.setUdfFuncTenantCodeMap(udfFuncMap);
        }
    }

    /**
     * get execute local path
     *
     * @return execute local path
     */
    private String getExecLocalPath(TaskInstance taskInstance) {
        return FileUtils.getProcessExecDir(taskInstance.getProcessDefine().getProjectId(),
                taskInstance.getProcessDefine().getId(),
                taskInstance.getProcessInstance().getId(),
                taskInstance.getId());
    }


    /**
     * whehter tenant is null
     *
     * @param tenant       tenant
     * @param taskInstance taskInstance
     * @return result
     */
    private boolean verifyTenantIsNull(String tenant, TaskInstance taskInstance) {
        if (tenant == null) {
            logger.error("tenant not exists,process instance id : {},task instance id : {}",
                    taskInstance.getProcessInstance().getId(),
                    taskInstance.getId());
            return true;
        }
        return false;
    }

    /**
     * get resource map key is full name and value is tenantCode
     */
    private Map<String, String> getResourceFullNames(TaskNode taskNode) {
        Map<String, String> resourceMap = new HashMap<>(16);
        AbstractParameters baseParam = TaskParametersUtils.getParameters(taskNode.getType(), taskNode.getParams());
        if (baseParam != null) {
            List<ResourceInfo> projectResourceFiles = baseParam.getResourceFilesList();
            if (CollectionUtils.isNotEmpty(projectResourceFiles)) {
                // filter the resources that the resource id equals 0
                Set<ResourceInfo> oldVersionResources = projectResourceFiles.stream()
                        .filter(t -> "".equals(t.getId())).collect(Collectors.toSet());
                if (CollectionUtils.isNotEmpty(oldVersionResources)) {
                    oldVersionResources.forEach((t) -> resourceMap.put(t.getRes(),
                            processService.queryTenantCodeByResName(t.getRes(), ResourceType.FILE))
                    );
                }

                // get the resource id in order to get the resource names in batch
                Stream<String> resourceIdStream = projectResourceFiles.stream().map(ResourceInfo::getId);
                Set<String> resourceIdsSet = resourceIdStream.collect(Collectors.toSet());

                if (CollectionUtils.isNotEmpty(resourceIdsSet)) {
                    String[] resourceIds = resourceIdsSet.toArray(new String[resourceIdsSet.size()]);

                    List<Resource> resources = processService.listResourceByIds(resourceIds);
                    resources.forEach(
                            (t) -> resourceMap.put(t.getFullName(), processService.queryTenantCodeByResName(t.getFullName(), ResourceType.FILE))
                    );
                }
            }
        }

        return resourceMap;
    }
}
